#!/usr/bin/env python
"""
启动自检脚本
用于校验端口、CORS和OLLAMA_BASE路径配置
"""
import os
import sys
import requests
import json
import time
from urllib.parse import urlparse

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入配置
from backend.common.config import config_manager
from backend.settings import HOST, PORT, CORS_ORIGINS

# 颜色定义
class Colors:
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    RED = '\033[91m'
    RESET = '\033[0m'

# 检查函数
class ConfigChecker:
    def __init__(self):
        # 获取配置
        self.host = HOST
        self.port = PORT
        self.cors_origins = CORS_ORIGINS
        self.ollama_config = config_manager.llm_providers.ollama
        self.ollama_base_url = self.ollama_config.base_url
        self.api_base_url = f"http://{self.host}:{self.port}"
        
    def check_port_availability(self):
        """检查端口可用性"""
        print(f"\n{Colors.YELLOW}检查端口 {self.port} 可用性...{Colors.RESET}")
        try:
            import socket
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(2)
                result = s.connect_ex((self.host, self.port))
                if result == 0:
                    print(f"{Colors.RED}端口 {self.port} 已被占用！请关闭占用该端口的程序或修改配置。{Colors.RESET}")
                    return False
                else:
                    print(f"{Colors.GREEN}端口 {self.port} 可用。{Colors.RESET}")
                    return True
        except Exception as e:
            print(f"{Colors.RED}端口检查失败: {str(e)}{Colors.RESET}")
            return False
    
    def check_ollama_connection(self):
        """检查OLLAMA服务连接"""
        print(f"\n{Colors.YELLOW}检查OLLAMA服务连接: {self.ollama_base_url}...{Colors.RESET}")
        try:
            # 尝试访问OLLAMA的tags端点
            tags_url = f"{self.ollama_base_url}/api/tags"
            response = requests.get(tags_url, timeout=5)
            
            if response.status_code == 200:
                try:
                    data = response.json()
                    models = data.get('models', [])
                    print(f"{Colors.GREEN}成功连接到OLLAMA服务！")
                    print(f"可用模型数量: {len(models)}")
                    if models:
                        print(f"前5个模型: {', '.join([m['name'] for m in models[:5]])}")
                    print(f"OLLAMA_BASE_URL格式正确: 不含/v1的根地址{Colors.RESET}")
                    return True
                except json.JSONDecodeError:
                    print(f"{Colors.YELLOW}成功连接到OLLAMA服务，但无法解析响应内容。{Colors.RESET}")
                    return True
            else:
                print(f"{Colors.RED}OLLAMA服务响应失败: HTTP {response.status_code}")
                print(f"响应内容: {response.text}")
                print(f"请确保OLLAMA服务正在运行，且URL格式正确。{Colors.RESET}")
                return False
        except requests.exceptions.ConnectionError:
            print(f"{Colors.RED}无法连接到OLLAMA服务！请确保OLLAMA服务正在运行。{Colors.RESET}")
            return False
        except Exception as e:
            print(f"{Colors.RED}OLLAMA服务检查失败: {str(e)}{Colors.RESET}")
            return False
    
    def check_cors_config(self):
        """检查CORS配置"""
        print(f"\n{Colors.YELLOW}检查CORS配置...{Colors.RESET}")
        expected_origins = ["http://localhost:5173", "http://127.0.0.1:5173"]
        
        print(f"当前CORS origins配置: {self.cors_origins}")
        
        missing_origins = [origin for origin in expected_origins if origin not in self.cors_origins]
        
        if not missing_origins:
            print(f"{Colors.GREEN}CORS配置正确，已包含所有必要的前端地址。{Colors.RESET}")
            return True
        else:
            print(f"{Colors.YELLOW}注意: CORS配置缺少以下前端地址: {', '.join(missing_origins)}")
            print(f"请在configs/config.yaml中的api.cors_origins添加这些地址。{Colors.RESET}")
            return False
    
    def check_api_endpoint(self):
        """检查API端点可用性（如果服务已启动）"""
        print(f"\n{Colors.YELLOW}尝试检查API端点可用性...{Colors.RESET}")
        try:
            health_url = f"{self.api_base_url}/v1/health"
            response = requests.get(health_url, timeout=2)
            
            if response.status_code == 200:
                print(f"{Colors.GREEN}API服务已启动并运行正常！{Colors.RESET}")
                return True
            else:
                print(f"{Colors.YELLOW}API服务可能正在启动中或返回了非200状态码: {response.status_code}{Colors.RESET}")
                return False
        except requests.exceptions.ConnectionError:
            print(f"{Colors.YELLOW}API服务尚未启动，请先启动服务后再次检查。{Colors.RESET}")
            return False
        except Exception as e:
            print(f"{Colors.RED}API端点检查失败: {str(e)}{Colors.RESET}")
            return False
    
    def run_full_check(self):
        """运行完整检查"""
        print(f"{Colors.YELLOW}=== 启动自检脚本 ==={Colors.RESET}")
        print(f"当前环境: {'开发环境' if os.environ.get('DEV_MODE') else '生产环境'}")
        
        # 检查配置有效性
        checks = [
            ('端口可用性', self.check_port_availability),
            ('OLLAMA连接', self.check_ollama_connection),
            ('CORS配置', self.check_cors_config)
        ]
        
        all_checks_passed = True
        for name, check_func in checks:
            if not check_func():
                all_checks_passed = False
        
        # 检查API端点（可选）
        self.check_api_endpoint()
        
        print(f"\n{Colors.YELLOW}=== 自检完成 ==={Colors.RESET}")
        
        if all_checks_passed:
            print(f"{Colors.GREEN}所有必要配置检查通过！可以正常启动服务。{Colors.RESET}")
        else:
            print(f"{Colors.RED}某些配置检查失败，请修复后再启动服务。{Colors.RESET}")
        
        return all_checks_passed

if __name__ == "__main__":
    checker = ConfigChecker()
    checker.run_full_check()